
#[derive(Debug, PartialEq, Eq, Clone)]
enum Packet {
    Integer(i32), // 21
    List(Vec<Packet>)  // [21, [21]]
}

impl Packet {
    fn from_str(input: &str) -> Packet {
        Packet::parse_str(input).0
    }

    fn parse_str(input: &str) -> (Packet, &str) {
        assert!(input.len() != 0);
        if input.starts_with("[]") {
            return (Packet::List(Vec::new()), &input[2..]);
        }

        if input.starts_with('[') {
            let mut packets = Vec::new();
            let mut input = &input[1..];
            loop {
                let (packet, left_input) = Packet::parse_str(input);
                packets.push(packet);
                match &left_input[0..1] {
                    "]" => return (Packet::List(packets), &left_input[1..]),
                    "," => input = &left_input[1..],
                    _ => panic!("Bad input"),
                }
            }
        } else {
            let digit_str = input
                .chars()
                .take_while(|c| c.is_digit(10))
                .collect::<String>();
            
            if digit_str.len() == 0 {
                panic!("Bad input!");
            }

            (Packet::Integer(digit_str.parse().unwrap()), &input[digit_str.len()..])
        }
    }
}

impl PartialOrd for Packet {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.compare(other))
    }
}

impl Ord for Packet {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.compare(other)
    }
}

impl Packet {
    fn compare(&self, other: &Packet) -> std::cmp::Ordering {
        match self {
            Packet::Integer(self_int) => {
                match other {
                    Packet::Integer(other_int) => {
                        self_int.cmp(other_int)
                    }
                    Packet::List(_) => {
                        Packet::List(vec![Packet::Integer(*self_int)]).compare(other)
                    }
                }
            }
            Packet::List(self_packets) => {
                match other {
                    Packet::Integer(other_int) => {
                        self.compare(&Packet::List(vec![Packet::Integer(*other_int)]))
                    }
                    Packet::List(other_packets) => {
                        for (p1, p2) in Iterator::zip(self_packets.iter(), other_packets.iter()) {
                            let res = p1.compare(p2);
                            match res {
                                std::cmp::Ordering::Less | std::cmp::Ordering::Greater => return res,
                                std::cmp::Ordering::Equal => {}
                            }
                        }

                        self_packets.len().cmp(&other_packets.len())
                    }
                }
            }
        }
    }
}

fn part1() {
    fn parse_input() -> Vec<(Packet, Packet)> {
        let input = include_str!("../input.txt");
        input
            .split("\r\n\r\n")
            .map(|two_lines| -> (Packet, Packet) {
                let mut iter = two_lines.lines();
                let line1 = iter.next().unwrap();
                let line2 = iter.next().unwrap();
                (Packet::from_str(line1), Packet::from_str(line2))
            })
            .collect()
    }

    let sum = parse_input()
        .into_iter()
        .enumerate()
        .map(|(i, (p_left, p_right))| (i + 1, p_left.cmp(&p_right)))
        .filter(|(_, res)| *res == std::cmp::Ordering::Less)
        .map(|(i, _)| i)
        .sum::<usize>();

    println!("Part 1: {}", sum);
}

fn part2() {
    fn parse_input() -> Vec<Packet> {
        include_str!("../input.txt")
            .lines()
            .filter(|line| !line.is_empty())
            .map(Packet::from_str)
            .collect()
    }

    let mut packets = parse_input();
    let divider_packet1 = Packet::List(vec![Packet::List(vec![Packet::Integer(2)])]);
    let divider_packet2 = Packet::List(vec![Packet::List(vec![Packet::Integer(6)])]);
    // Add two additional divider packets:
    packets.push(divider_packet1.clone());
    packets.push(divider_packet2.clone());

    packets.sort();

    
    let p = packets
        .into_iter()
        .enumerate()
        .filter(|(_, p)| p == &divider_packet1 || p == &divider_packet2)
        .map(|(i, _)| i + 1)
        .product::<usize>();
    
    println!("Part 2: {}", p);
}

fn main() {
    part1();
    part2();
}

#[cfg(test)]
mod test {

    use super::*;

    #[test]
    fn test_create_packet() {
        assert_eq!(
            Packet::from_str("[]"),
            Packet::List(Vec::new()),
        );

        assert_eq!(
            Packet::from_str("12"),
            Packet::Integer(12),
        );

        assert_eq!(
            Packet::from_str("[1]"),
            Packet::List(vec![Packet::Integer(1)])
        );

        assert_eq!(
            Packet::from_str("[1,2,3]"),
            Packet::List(vec![Packet::Integer(1), Packet::Integer(2), Packet::Integer(3)])
        );

        assert_eq!(
            Packet::from_str("[[[],[],8,3],[10]]"),
            Packet::List(vec![
                Packet::List(vec![
                    Packet::List(vec![]),
                    Packet::List(vec![]),
                    Packet::Integer(8),
                    Packet::Integer(3),
                ]),
                Packet::List(vec![
                    Packet::Integer(10),
                ]),
            ])
        );

    }

    #[test]
    fn test_cmp() {
        assert_eq!(
            Packet::from_str("[1,1,3,1,1]").cmp(&Packet::from_str("[1,1,5,1,1]")),
            std::cmp::Ordering::Less,
        );

        assert_eq!(
            Packet::from_str("[[1],[2,3,4]]").cmp(&Packet::from_str("[[1],4]")),
            std::cmp::Ordering::Less,
        );

        assert_eq!(
            Packet::from_str("[9]").cmp(&Packet::from_str("[[8,7,6]]")),
            std::cmp::Ordering::Greater,
        );

        assert_eq!(
            Packet::from_str("[[4,4],4,4]").cmp(&Packet::from_str("[[4,4],4,4,4]")),
            std::cmp::Ordering::Less,
        );

        assert_eq!(
            Packet::from_str("[7,7,7,7]").cmp(&Packet::from_str("[7,7,7]")),
            std::cmp::Ordering::Greater,
        );

        assert_eq!(
            Packet::from_str("[]").cmp(&Packet::from_str("[3]")),
            std::cmp::Ordering::Less,
        );

        assert_eq!(
            Packet::from_str("[[[]]]").cmp(&Packet::from_str("[[]]")),
            std::cmp::Ordering::Greater,
        );

        assert_eq!(
            Packet::from_str("[1,[2,[3,[4,[5,6,7]]]],8,9]").cmp(&Packet::from_str("[1,[2,[3,[4,[5,6,0]]]],8,9]")),
            std::cmp::Ordering::Greater,
        );
    }
}